# src/plot_data/plotamglevel.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*

import os
import json
import argparse
import sys 
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt


sys.path.append("..")
from package.utility import ensure_exists
from package.utility import find_targets
from baseplot import plot_linechart
from baseplot import plot_percent_2D
from baseplot import plot_percent
from baseplot import plot_percent_line
from baseplot import plot_heatmap
from baseplot import casename2number
from baseplot import getplteffdata
from baseplot import create_subplot
from baseplot import close_sublot_levelplot
from baseplot import mypaper_subplot_level
from baseplot import mypaper_subplot_level_axset

# set parser
def make_parser(): 
    parser = argparse.ArgumentParser(
        prog="jxpamg-predo-merge-pid-timerdata",
        description="jxpamg timer data preprocess script, merge pid's timer data to a single json.")
    parser.add_argument(
            "-sdir","--savedir",
            dest="savedir",
            default="./",
            help="Destination save pic/csv directories."
        )
    parser.add_argument(
            "-o", "--output",
            dest="out_type",
            default="csv",
            help="Set output type of plotting picture. Option: csv, pic."
    )
    parser.add_argument(
            "-po", "--picopt",
            dest="plot_type",
            default="time_linechart",
            help="Set type of pic out. Option: time_linechart."
    )
    parser.add_argument(
            "-co", "--csvopt",
            dest="csv_type",
            default="time",
            help="Set type of data in csv. Option: time."
    )
    parser.add_argument(
            "-np", "--numprocess",
            dest="numprocess",
            default=32,
            help="Set number of process for plotting"
    )
    parser.add_argument(
            "LOGDIR",
            metavar="LOGDIR",
            help="Destination log directories."
        )
    parser.add_argument(
            "COMPONENT",
            metavar="COMPONENT",
            default="CreateS",
            help="Set component of plotting picture. Option: ."
    )
    return parser

def get_runcases(topdir,jsonname="*-database.json"):
    """Get runcase under topdir."""
    topdir = os.path.abspath(topdir)
    ensure_exists(topdir)

    print(f">> scan runcase under {topdir}") 
    runcases = find_targets(topdir, jsonname, suffix="-xtype f") 
    runcase_num = len(runcases)
    print(f">> {runcase_num} runcases found!")

    return runcases

def gettimelevel(nump,dict,result_list):
    # TimeLevel = [0] * (dict["AmgLevel"] + 1)
    if "AmgLevel" in dict and "AmgCycle" not in dict:
        for k in dict["TimeLevel"]:
            resulr_dict = {}
            resulr_dict["np"] = nump
            level = int(k.split("_")[-1])
            resulr_dict["level"] = level
            resulr_dict["time"] = float(dict["TimeLevel"][k]["tavg"])
            result_list.append(resulr_dict)
    elif "AmgLevel" in dict and "AmgCycle" in dict:
        Tag = 0
        tmp_list = []
        for k in dict["TimeLevel"]:
            for l in dict["TimeLevel"][k]:
                level = int(l.split("_")[-1])
                if Tag == 0:
                    resulr_dict = {}
                    resulr_dict["np"] = nump
                    resulr_dict["level"] = level
                    resulr_dict["time"] = float(dict["TimeLevel"][k][l]["tavg"])
                    tmp_list.append(resulr_dict)
                else:
                    #TODO:下面应该可以改的更简单，先用数组存起来都行
                    for tmpdict in tmp_list:
                        # print(tmpdict)
                        if tmpdict["level"] == level:
                            tmpdict["time"] += float(dict["TimeLevel"][k][l]["tavg"])
                        # print(tmp_list)
            Tag += 1
        for i in tmp_list:
            result_list.append(i)
    # print(result_list)
    return result_list

# 获取对应 level 和 parent 的 模块 的 ave 时间
def getavetime(jsondata, component, plttype):
# TODO:这里写的更多样化一点
# 针对不同画图类型有不同的参数提取
    datalist = []
    for i in jsondata:
        # 加入np
        nump = int(i.split("_")[1])
        if(plttype == "percent_line"):
            for j in jsondata[i]:
                if j == component :
                    datalist = gettimelevel(nump, jsondata[i][j],datalist)
        else:
            #paren == NULL,取得level相同的组件
            for j in jsondata[i]:
                if j == component :
                    datalist = gettimelevel(nump, jsondata[i][j],datalist)
        # print(tempdic)
    # print(datalist)
    df = pd.DataFrame(datalist)
    return df

# 获得不同 case 得到 component 对应的 time
def getjsondata(casename, component, plttype):
    timedic = {}
    for i in casename:
        basename = os.path.basename(i)
        basename = basename.split("-")[0]
        with open(i, "r") as fo:
            text = fo.read()
            jsondata = json.loads(text)
            print(f"Load {basename} data.")
        timedic[basename] = getavetime(jsondata, component, plttype)
    #print("timedic:",timedic)
    return timedic

def plotdict(numprocess,dict,savedir):
    # print(dict)
    # 画出每一个case，指定某一进程，每层折线图。
    # mylen = 0
    #print("dict:",dict)
    legend_name = []
    x = []
    y = []
    for i in dict:
        legend_name.append(i)
        for j in dict[i]:
            # j是一个字典
            x.append(j["np"])
            y.append(j["Totaltime"])
            if(j["np"] == numprocess):
                plt.plot(j["TimeLevellist"],j["TimeLevel"],marker = 'o')
    legend_name = casename2number(legend_name)
    plt.xticks(fontsize = 8)
    plt.yticks(fontsize = 8)
    plt.xlabel(r"Level",fontsize = 16)
    plt.ylabel(r"Time(s)",fontsize = 16)
    title = j["Tag"] + "-TimeLevel"+"-np-" + f"{numprocess}"
    plt.title(title,fontsize = 20)
    plt.legend(legend_name,fontsize = 6)
    savedir1 = os.path.join(savedir, title +"-linechart.png")
    plt.savefig(savedir1)
    plt.close()
                # 下面两个数组分别画图或者生成一个numpy，但是每个算例的层不一样
                # 建议下面直接画图
                # print(i)
                # # print(j["TimeLevellist"])
                # print(j["TimeLevel"])
                # print("\n\n")
                #mylen = max(len(j["TimeLevel"]), mylen)

    
    return

def plotdict_2nd(dict,savedir):
    # print(dict)
    # 画出每一个case，指定某一进程，每层折线图。
    #print("dict:",dict)
    for i in dict:
        legend_name = []
        mylen = 10000
        y = []
        xxx = []
        for j in dict[i]:
            mylen = min(mylen,len(j["TimeLevel"]))
        xx = range(len(dict[i]))
        for j in dict[i]:
            # j是一个字典
            xxx.append(j["np"])
            y = np.append(y,j["np"])
            y = np.append(y,j["TimeLevel"][0:mylen])
        legend_name = j["TimeLevellist"][0:mylen]
        xxx.sort()
        #print(y)
        #print(mylen)
        #print(len(y))
        y.resize(len(dict[i]),mylen+1)
        x = y[:,0]
        index = np.lexsort((x,))
        y = y[index]
        id = casename2number([i])
        #print(id)
        if(len(j["Tag"].split("-")) > 1):
            title =  f"c{id[0]}-" + j["Tag"].split("-")[1]+ "-LevelTime"
        else:
            title =  f"c{id[0]}-" + j["Tag"]+ "-LevelTime"
        # print(i)
        #print(dict[i])
        for j in range(1,mylen+1,1):
            plt.plot(xx,y[:,j],marker = 'o')
            #print(j)
        # print(y)
        # print(len(y))
    # #     plt.plot(x,y,marker = 'o')
        #legend_name = casename2number(legend_name)
        plt.xticks(xx,xxx)
        plt.xticks(fontsize = 8)
        plt.yticks(fontsize = 8)
        plt.xlabel(r"num_proc",fontsize = 16)
        plt.ylabel(r"Time(s)",fontsize = 16)
        plt.title(title,fontsize = 20)
        plt.legend(legend_name,fontsize = 6)
        savedir1 = os.path.join(savedir, title +"-linechart.png")
        plt.savefig(savedir1)
        plt.close()

def mysubplot(dict,component,savedir):
    legend_name = []
    fig, axes = create_subplot(nrows=2, ncols=3, xlable="Number  of  Processors", ylable="Time(s)")

    for i in dict:
        caseid = casename2number(i)
        casenumber = int(caseid.split("C")[-1]) - 1
        df = dict[i]
        levellist = list(df["level"])
        levellist = sorted(set(levellist), key=levellist.index)
        title = f"{caseid}_Level-{component}"
        for j in levellist:
            if j not in legend_name:
                legend_name.append(f"Level {j}")
            levelplotdf = df[df["level"]==j].sort_values(by=["np"])
            levelplotdf = levelplotdf.drop(["level"], axis=1)
            levelplotdf.set_index(["np"], inplace=True)
            mypaper_subplot_level(axes[casenumber//3, casenumber%3], levelplotdf)
        mypaper_subplot_level_axset(axes[casenumber//3, casenumber%3], title, levelplotdf)
    close_sublot_levelplot(2,3,fig,axes,savedir,title,"leveltime",legend_name)

def main():
    parser = make_parser()
    args = parser.parse_args()
    runcases = get_runcases(args.LOGDIR,"*-database_2nd.json")
    datadict = getjsondata(runcases, args.COMPONENT, args.plot_type)
    
    #TODO:plot
    # plotdict(args.numprocess, datadict,args.savedir)
    # plotdict_2nd(datadict,args.savedir)
    print(datadict)
    mysubplot(datadict,args.COMPONENT,args.savedir)

if __name__ == "__main__":
    main()